#-*- coding:utf-8 -*-
# author:贤宁
# datetime:2021/11/12 10:07
# software: PyCharm

import torch
import torch.nn as nn
from torch.nn import init

class ConvGRUCell(nn.Module):
    def __init__(self, input_size, hidden_size, kernel_size, activation=torch.sigmoid):
        super(ConvGRUCell, self).__init__()
        padding = kernel_size // 2
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.reset_gate = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size, padding=padding, stride=1)
        self.update_gate = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size, padding=padding, stride=1)
        self.out_gate = nn.Conv2d(input_size + hidden_size, hidden_size, kernel_size, padding=padding, stride=1)
        self.activation = activation

        init.orthogonal_(self.reset_gate.weight)
        init.orthogonal_(self.update_gate.weight)
        init.orthogonal_(self.out_gate.weight)
        init.constant_(self.reset_gate.bias, 0.)
        init.constant_(self.update_gate.bias, 0.)
        init.constant_(self.out_gate.bias, 0.)

    def forward(self, x, prev_state=None):
        output = []
        for step in range(x.size(1)):
            combined_1 = torch.cat([x[:, step, :, :, :], prev_state], dim=1)
            #print(combined_1.size())
            update = self.activation(self.update_gate(combined_1))
            #print(update.size())
            reset = self.activation(self.reset_gate(combined_1))
           # print(reset.size())
            out_inputs = torch.tanh(self.out_gate(torch.cat([x[:, step, :, :, :], prev_state * reset], dim=1)))
            new_state = prev_state * (1 - update) + out_inputs * update
            output.append(new_state)
        output = torch.stack(output, 0)
        #print(output.size())
        return output